import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='3'
import cv2
import matplotlib
import matplotlib.pyplot as plt
matplotlib.use("Agg")
import numpy as np
import tensorflow as tf
tf.get_logger().setLevel('ERROR')
from know import KNOW
from know import MEASURE
from know import PARAM

SIZE,DIM,TARGET=1,50,64
ZOOM,GROUND=1,1200.0

_know=KNOW(36,50.0,1.0,10,50,150.0,15.0,5.0)

path=os.path.join("datasetsedd","test")
fs=[fs for parent,dirnames,fs in os.walk(path)][0]
fs=[os.path.join(path,det) for det in fs if ".dat" in det]

decoder=tf.keras.models.load_model("decoder.keras")

for each in fs[11:12]:
    print(each)
    basename=os.path.splitext(os.path.basename(each))[0]
    name=os.path.splitext(each)[0]
    content=[]
    with open(name+".res","r") as fp:
        # content=np.array([float(i) for i in [det.split() for det in fp][0]]).astype(np.float32)
        content=np.array([float(i) for i in fp]).astype(np.float32)
    content=tf.reshape(content,[-1,content.shape[0]])

    data=np.genfromtxt(each)

    if os.path.exists(name) is False:
        os.mkdir(name)
    
    for i in range(SIZE):
        # z=tf.random.normal(shape=(1,DIM),mean=0.0,stddev=1.0)
        # z_lbl_concat=np.concatenate((z,content),axis=1)
        # preds=decoder.predict(z_lbl_concat)
        # np.savetxt(os.path.join(name,"result.pred"),preds.reshape((TARGET,TARGET)))

        preds=np.genfromtxt(os.path.join(name,"result.pred"))
        
        generated_digit=tf.reshape(preds,[TARGET,TARGET])
        # generated_digit=(generated_digit*std)+mean
        generated_digit=generated_digit*255.
        image=generated_digit.numpy()

        plt.figure()
        plt.imshow(image,cmap='gray')
        plt.axis('Off')
        plt.savefig(os.path.join(name,"result.jpg"))
        plt.cla()
        plt.clf()
        plt.close()

        difference=np.absolute(data-preds.reshape((TARGET,TARGET)))
        image=difference*255.        
        np.savetxt(os.path.join(name,"result_.pred"),image.reshape((TARGET,TARGET)).astype(np.int32))

        plt.figure()
        plt.imshow(image,cmap='gray')
        plt.axis('Off')
        plt.savefig(os.path.join(name,"result_.jpg"))
        plt.cla()
        plt.clf()
        plt.close()

        preds=preds.reshape((TARGET,TARGET))
        nx,ny=_know.XRANGE,_know.NLAYER
        top_border,bottom_border,left_border,right_border=int((TARGET-ny*ZOOM)/2.),int((TARGET-ny*ZOOM)/2.),int((TARGET-nx*ZOOM)/2.),int((TARGET-nx*ZOOM)/2.)
        preds=preds[top_border:TARGET-bottom_border,left_border:TARGET-right_border]

        # preds=preds.flatten("C")
        # preds=preds.ravel()
        # array=[]
        # for det in preds:
        #     temp=abs(1.0-det)
        #     if temp<0.0:
        #         temp=1.0-abs(temp)
        #     array.append(np.log(temp*GROUND))
        # preds=cv2.resize(preds,(nx,ny))

        # for x in range(0,preds.shape[0]):
        #     for y in range(0,preds.shape[1]):
        #         temp=abs(1.0-preds[x,y])
        #         # if temp<0.0:
        #         #     temp=1.0-abs(temp)
        #         preds[x,y]=np.log(temp*GROUND)

        preds=np.log((abs(1.0-cv2.resize(preds,(nx,ny)))*GROUND))

        # np.nan_to_num(preds,False,0.0)
        # print(np.isnan(preds).any())
        
        with open(os.path.join(name,str(i)+".mod"),"w") as fp:        
            fp.write(" ".join([str(det) for det in [_know.XRANGE,_know.NLAYER,"LOGE"]])+"\n")
            fp.write(" ".join([str(det) for det in _know.INTERVAL])+"\n")
            fp.write(" ".join([str(det) for det in _know.TIC])+"\n")
            np.savetxt(fp,preds)



